# 添加节点
import time

import networkx as nx
import matplotlib.pyplot as plt
import copy
from create_alarms_data_002 import ALL_ORIGIN_INFOS

def draw(batch_data, data, epoch, idx):
    # print("ALL_ORIGIN_INFOS",ALL_ORIGIN_INFOS)
    G = copy.deepcopy(nx.Graph())
    nid = batch_data.n_id.tolist()
    # print(nid)
    # 0 -100 100-286 286-297
    # print(data)
    # print("This batch data is  ",batch_data)
    # ##############################
    # print(batch_data.edge_index)
    # print(batch_data.edge_type)
    #
    # print(batch_data.n_id)
    # print(batch_data.node_type)
    # print(batch_data.x.shape)
    # print(batch_data.edge_label_index)
    # print(batch_data.edge_label)

    """
    本次抽样数据信息
    Data(
     edge_index=[2, 67], 本次抽样有67条边
     node_id=[297], 一共参与节点有297个
     x=[45, 768],  本次抽样有 45 个点的特征 id列表用 n_id 表示。
     node_type=[45], 本次抽样点的类型。
     edge_type=[67], 本次抽样边的类型。
     edge_label=[3], 本次抽样中minibatch的边的标签
     edge_label_index=[2, 3], 本次抽样中minibatch的对应关系，该对应关系应该从n_id中找寻，另外，该数据是从edge_index的67条边中抽样出来的
     n_id=[45], 本次抽样有 45 个点 id列表用 n_id 表示。 ######## 节点id来自于 node_id 是从297个节点中抽样出45个点的结果。
     e_id=[67], 本次抽样有67条边
     input_id=[3])
    #batch_data.edge_index # 边的对应关系 有67条边
    tensor([[ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15,  4, 16, 17,  0, 18, 19,  2, 16,
          1,  0, 18, 20, 21, 22, 23, 18, 24,  0,  0, 18, 25, 26, 18, 27,  0, 14,
          1,  1, 28, 29, 30, 11,  1, 15, 31, 32,  1, 33, 34, 35,  2, 36, 18, 37,
         38, 39,  0, 17, 40, 41, 25, 42,  2, 43,  4, 36, 44],
        [ 0,  0,  0,  0,  0,  1,  1,  1,  1,  1,  2,  2,  3,  3,  3,  4,  4,  4,
          5,  6,  6,  6,  6,  7,  7,  7,  7,  7,  8,  8,  8,  9,  9, 10, 10, 11,
         11, 12, 12, 13, 13, 14, 14, 14, 14, 14, 15, 16, 16, 16, 16, 16, 17, 17,
         17, 17, 17, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19]])
    #batch_data.edge_type # 
    tensor([2, 2, 2, 2, 1, 1, 1, 2, 2, 1, 0, 5, 3, 4, 4, 4, 3, 4, 4, 5, 5, 0, 0, 0,
        0, 5, 0, 5, 5, 5, 0, 0, 5, 3, 4, 3, 4, 4, 3, 0, 0, 0, 5, 0, 0, 0, 4, 2,
        2, 2, 2, 1, 5, 0, 0, 0, 5, 2, 1, 1, 1, 2, 2, 2, 1, 1, 1])
    #batch_data.n_id 抽样中45个点和点的id值。
    tensor([  2,  10, 313, 453, 463, 466, 344, 221, 316, 220, 413, 398, 394, 324,
        200, 432,   5, 251,   0,   9, 428, 386, 418, 373, 450, 383, 423, 226,
        310, 470, 377, 468, 406, 317, 322, 286, 445, 420, 376, 467, 435, 403,
        306, 271, 400])
    #batch_data.x.shape 表明有45个节点参与了此次抽样。
    torch.Size([45, 768])
    #batch_data.edge_label_index
    tensor([[1, 4, 3],
        [5, 2, 0]])
    #batch_data.edge_label
    tensor([1, 1, 1])
    """

    #batch_data.n_id
    #batch_data.node_type
    #batch_data.edge_label_index
    id_list1 = [nid[i] for i in range(len(nid)) if batch_data.node_type.tolist()[i] == 0]
    id_list2 = [nid[i] for i in range(len(nid)) if batch_data.node_type.tolist()[i] == 1]
    id_list3 = [nid[i] for i in range(len(nid)) if batch_data.node_type.tolist()[i] == 2]

    G.add_nodes_from(id_list1, color='blue', size=5)
    G.add_nodes_from(id_list2, color='green', size=5)
    G.add_nodes_from(id_list3, color='orange', size=5)

    edge_true_mask = batch_data.edge_label == 1
    edge_false_mask = batch_data.edge_label == 0

    true_edges = batch_data.edge_label_index.t()[edge_true_mask].tolist()
    false_edges = batch_data.edge_label_index.t()[edge_false_mask].tolist()

    true_edge_list = [[nid[e[0]], nid[e[1]]] for e in true_edges]
    false_edge_list = [[nid[e[0]], nid[e[1]]] for e in false_edges]
    # print("1_nodes", nodes)
    # print("true_edge_list",true_edge_list)
    # 10.29.130.19|出现 CRITICAL DISK|陕西西安移动|LIVE_EPG_FCACHE
    # for true_edge in true_edge_list:
    #     print("start:{} {}".format(true_edge[0],ALL_ORIGIN_INFOS[true_edge[0]]))
    #     print("end:{} {}".format(true_edge[1],ALL_ORIGIN_INFOS[true_edge[1]]))

    # print("2_true_edge_list", true_edge_list)
    # print("3_false_edge_list", false_edge_list)
    #
    # for i in true_edge_list:
    #     print("4_true_like_data: {}".format(data["all_dict"][i[0]]+"||||"+data["all_dict"][i[1]]))

    G.add_edges_from(true_edge_list, color='red', size=10)
    G.add_edges_from(false_edge_list, color='gray', size=10)

    pos = nx.spring_layout(G, k=1, scale=5)
    # pos =  nx.spring_layout(G,k=2,scale=2)
    # pos = nx.circular_layout(G,scale=1000)
    # pos = nx.shell_layout(G)
    # pos = nx.spectral_layout(G)
    # print("##1111",[n[1]["shape"] if "shape" in n[1] else "8" for n in G.nodes(data=True)])

    nx.draw(G, pos=pos, with_labels=True, node_size=200,
            # style="--",
            # 'so^>v<dph8'
            node_shape="o",
            # font_size=[n[1]["size"] if "size" in n[1] else 10 for n in G.nodes(data=True)],
            font_size=10,
            node_color=[n[1]["color"] if "color" in n[1] else "black" for n in G.nodes(data=True)],
            edge_color=[n[2]["color"] if "color" in n[2] else "black" for n in G.edges(data=True)]
            # width =2.0
            )
    plt.show()
    # time.sleep(1)
    plt.savefig("/home/Dyf/code/storage_models/logs/output_{}_{}.png".format(epoch, idx))
    plt.close()